"""
数据源连接服务
支持多种数据库类型的数据源连接和查询
"""
import asyncio
from typing import Dict, Any, List, Optional, Union
from sqlalchemy import create_engine, text, inspect
from sqlalchemy.exc import SQLAlchemyError, OperationalError
import pandas as pd
from contextlib import contextmanager
import json
import time
from datetime import datetime

from core.config import settings
from core.logging_config import get_logger
from core.exceptions import DatabaseError, ValidationError

logger = get_logger("data_source_connection")


class DatabaseType:
    """数据库类型枚举"""
    MYSQL = "mysql"
    POSTGRESQL = "postgresql"
    SQLITE = "sqlite"
    ORACLE = "oracle"
    SQLSERVER = "sqlserver"


class DataSourceConnection:
    """数据源连接管理类"""
    
    def __init__(self, connection_config: Dict[str, Any]):
        """
        初始化数据源连接
        
        Args:
            connection_config: 连接配置
                - type: 数据库类型
                - host: 主机地址
                - port: 端口号
                - database: 数据库名
                - username: 用户名
                - password: 密码
                - extra_params: 额外连接参数
        """
        self.config = self._normalize_config(connection_config)
        self.engine = None
        self.connection_info = None
        self.last_test_time = None
        self.is_connected = False
    
    def _normalize_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """标准化配置字典，处理字段名映射"""
        normalized = config.copy()
        
        # 处理数据库类型字段名映射
        if "database_type" in normalized:
            normalized["type"] = normalized.pop("database_type")
        elif "db_type" in normalized:
            normalized["type"] = normalized.pop("db_type")
        
        # 确保必需字段存在
        if "extra_params" not in normalized:
            normalized["extra_params"] = {}
        
        return normalized
        
    def _build_connection_string(self) -> str:
        """构建数据库连接字符串"""
        db_type = self.config.get("type") or self.config.get("database_type") or self.config.get("db_type", "")
        if not db_type:
            raise ValidationError("数据库类型不能为空")
        
        db_type = db_type.lower()
        
        if db_type == DatabaseType.MYSQL:
            return self._build_mysql_connection()
        elif db_type == DatabaseType.POSTGRESQL:
            return self._build_postgresql_connection()
        elif db_type == DatabaseType.SQLITE:
            return self._build_sqlite_connection()
        elif db_type == DatabaseType.ORACLE:
            return self._build_oracle_connection()
        elif db_type == DatabaseType.SQLSERVER:
            return self._build_sqlserver_connection()
        else:
            raise ValidationError(f"不支持的数据库类型: {db_type}")
    
    def _build_mysql_connection(self) -> str:
        """构建MySQL连接字符串"""
        host = self.config.get("host", "localhost")
        port = self.config.get("port", 3306)
        database = self.config.get("database", "")
        username = self.config.get("username", "")
        password = self.config.get("password", "")
        extra_params = self.config.get("extra_params", {})
        
        # 验证必需字段
        if not database:
            raise ValidationError("数据库名不能为空")
        if not username:
            raise ValidationError("用户名不能为空")
        
        # 处理密码中的特殊字符
        from urllib.parse import quote_plus
        encoded_password = quote_plus(password) if password else ""
        
        # 构建额外参数
        params = []
        for key, value in extra_params.items():
            params.append(f"{key}={value}")
        
        params_str = "&".join(params) if params else ""
        
        if params_str:
            connection_string = f"mysql+pymysql://{username}:{encoded_password}@{host}:{port}/{database}?{params_str}"
        else:
            connection_string = f"mysql+pymysql://{username}:{encoded_password}@{host}:{port}/{database}"
        
        return connection_string
    
    def _build_postgresql_connection(self) -> str:
        """构建PostgreSQL连接字符串"""
        host = self.config.get("host", "localhost")
        port = self.config.get("port", 5432)
        database = self.config.get("database", "")
        username = self.config.get("username", "")
        password = self.config.get("password", "")
        extra_params = self.config.get("extra_params", {})
        
        # 验证必需字段
        if not database:
            raise ValidationError("数据库名不能为空")
        if not username:
            raise ValidationError("用户名不能为空")
        
        # 处理密码中的特殊字符
        from urllib.parse import quote_plus
        encoded_password = quote_plus(password) if password else ""
        
        # 构建额外参数
        params = []
        for key, value in extra_params.items():
            params.append(f"{key}={value}")
        
        params_str = "&".join(params) if params else ""
        
        if params_str:
            connection_string = f"postgresql+psycopg2://{username}:{encoded_password}@{host}:{port}/{database}?{params_str}"
        else:
            connection_string = f"postgresql+psycopg2://{username}:{encoded_password}@{host}:{port}/{database}"
        
        return connection_string
    
    def _build_sqlite_connection(self) -> str:
        """构建SQLite连接字符串"""
        database = self.config.get("database", "")
        if not database:
            raise ValidationError("SQLite数据库路径不能为空")
        
        return f"sqlite:///{database}"
    
    def _build_oracle_connection(self) -> str:
        """构建Oracle连接字符串"""
        host = self.config.get("host", "localhost")
        port = self.config.get("port", 1521)
        database = self.config.get("database", "")
        username = self.config.get("username", "")
        password = self.config.get("password", "")
        
        connection_string = f"oracle+cx_oracle://{username}:{password}@{host}:{port}/{database}"
        return connection_string
    
    def _build_sqlserver_connection(self) -> str:
        """构建SQL Server连接字符串"""
        host = self.config.get("host", "localhost")
        port = self.config.get("port", 1433)
        database = self.config.get("database", "")
        username = self.config.get("username", "")
        password = self.config.get("password", "")
        
        connection_string = f"mssql+pymssql://{username}:{password}@{host}:{port}/{database}"
        return connection_string
    
    def test_connection(self) -> Dict[str, Any]:
        """测试数据库连接"""
        try:
            connection_string = self._build_connection_string()
            
            # 记录连接字符串（隐藏密码）
            safe_connection_string = connection_string
            if '@' in safe_connection_string:
                safe_connection_string = safe_connection_string.split('@')[0].split(':')[0] + ':***@' + safe_connection_string.split('@')[1]
            logger.info(f"构建的连接字符串: {safe_connection_string}")
            
            # 创建引擎
            engine = create_engine(
                connection_string,
                pool_pre_ping=True,
                pool_recycle=3600
            )
            
            # 测试连接
            start_time = time.time()
            with engine.connect() as conn:
                # 执行简单查询
                result = conn.execute(text("SELECT 1"))
                result.fetchone()
            
            connection_time = time.time() - start_time
            
            # 获取数据库信息
            inspector = inspect(engine)
            db_info = {
                "version": inspector.get_table_names().__class__.__module__,
                "table_count": len(inspector.get_table_names()),
                "connection_time": connection_time
            }
            
            self.engine = engine
            self.connection_info = db_info
            self.last_test_time = datetime.utcnow()
            self.is_connected = True
            
            logger.info(f"数据库连接测试成功: {self.config.get('type')}")
            
            return {
                "success": True,
                "message": "连接测试成功",
                "connection_time": connection_time,
                "database_info": db_info
            }
            
        except Exception as e:
            self.is_connected = False
            logger.error(f"数据库连接测试失败: {str(e)}")
            return {
                "success": False,
                "message": f"连接测试失败: {str(e)}",
                "connection_time": None,
                "database_info": None
            }
    
    def get_table_list(self) -> List[Dict[str, Any]]:
        """获取数据库表列表"""
        if not self.is_connected or not self.engine:
            raise DatabaseError("数据库未连接")
        
        try:
            inspector = inspect(self.engine)
            tables = []
            
            for table_name in inspector.get_table_names():
                columns = []
                for column in inspector.get_columns(table_name):
                    columns.append({
                        "name": column["name"],
                        "type": str(column["type"]),
                        "nullable": column.get("nullable", True),
                        "default": column.get("default"),
                        "primary_key": column.get("primary_key", False)
                    })
                
                tables.append({
                    "name": table_name,
                    "columns": columns,
                    "column_count": len(columns)
                })
            
            return tables
            
        except Exception as e:
            logger.error(f"获取表列表失败: {str(e)}")
            raise DatabaseError(f"获取表列表失败: {str(e)}")
    
    def execute_query(self, query: str, limit: int = 1000) -> Dict[str, Any]:
        """执行SQL查询"""
        if not self.is_connected or not self.engine:
            raise DatabaseError("数据库未连接")
        
        try:
            start_time = time.time()
            
            # 使用pandas执行查询
            with self.engine.connect() as conn:
                df = pd.read_sql_query(query, conn)
            
            # 限制返回行数
            if len(df) > limit:
                df = df.head(limit)
            
            # 转换为字典格式
            data = df.to_dict('records')
            
            # 获取列信息
            columns = []
            for col in df.columns:
                columns.append({
                    "name": col,
                    "type": str(df[col].dtype),
                    "sample_values": df[col].head(3).tolist()
                })
            
            execution_time = time.time() - start_time
            
            return {
                "success": True,
                "data": data,
                "columns": columns,
                "row_count": len(data),
                "execution_time": execution_time,
                "total_rows": len(df)
            }
            
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}")
            return {
                "success": False,
                "message": f"查询执行失败: {str(e)}",
                "data": [],
                "columns": [],
                "row_count": 0,
                "execution_time": None,
                "total_rows": 0
            }
    
    def get_table_data(self, table_name: str, limit: int = 10000) -> Dict[str, Any]:
        """获取表数据"""
        query = f"SELECT * FROM {table_name} LIMIT {limit}"
        return self.execute_query(query, limit)
    
    def preview_table_data(self, table_name: Union[str, Dict[str, Any]], limit: int = 100, offset: int = 0) -> Dict[str, Any]:
        """预览表数据（支持分页）"""
        try:
            # 处理表名参数，支持字符串或字典
            if isinstance(table_name, dict):
                if 'name' in table_name:
                    actual_table_name = table_name['name']
                else:
                    raise ValidationError("表名字典缺少 'name' 字段")
            else:
                actual_table_name = table_name
            
            # 验证表名
            if not actual_table_name or not isinstance(actual_table_name, str):
                raise ValidationError("表名不能为空且必须是字符串")
            
            # 防止SQL注入：验证表名只包含安全字符
            import re
            if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', actual_table_name):
                raise ValidationError(f"表名包含非法字符: {actual_table_name}")
            
            # 构建分页查询
            query = f"SELECT * FROM {actual_table_name} LIMIT {limit} OFFSET {offset}"
            
            # 获取总记录数
            count_query = f"SELECT COUNT(*) as total FROM {actual_table_name}"
            
            start_time = time.time()
            
            with self.engine.connect() as conn:
                # 执行计数查询
                count_result = conn.execute(text(count_query))
                total_records = count_result.fetchone()[0]
                
                # 执行数据查询
                df = pd.read_sql_query(query, conn)
                
            # 转换为字典格式
            data = df.to_dict('records')
            
            # 获取列信息
            columns = []
            for col in df.columns:
                columns.append({
                    "name": col,
                    "type": str(df[col].dtype),
                    "sample_values": df[col].head(3).tolist() if len(df) > 0 else []
                })
            
            execution_time = int((time.time() - start_time) * 1000)
            
            return {
                "table_name": actual_table_name,
                "columns": columns,
                "data": data,
                "total_records": total_records,
                "returned_records": len(data),
                "offset": offset,
                "limit": limit,
                "execution_time": execution_time,
                "column_count": len(columns)
            }
            
        except Exception as e:
            logger.error(f"预览表数据失败: {str(e)}")
            # 在错误情况下，尝试获取原始表名
            error_table_name = actual_table_name if 'actual_table_name' in locals() else (table_name['name'] if isinstance(table_name, dict) else str(table_name))
            
            return {
                "table_name": error_table_name,
                "columns": [],
                "data": [],
                "total_records": 0,
                "returned_records": 0,
                "offset": offset,
                "limit": limit,
                "execution_time": 0,
                "column_count": 0,
                "error": str(e)
            }
    
    def get_table_schema(self, table_name: str) -> Dict[str, Any]:
        """获取表结构"""
        if not self.is_connected or not self.engine:
            raise DatabaseError("数据库未连接")
        
        try:
            inspector = inspect(self.engine)
            
            # 获取列信息
            columns = []
            for column in inspector.get_columns(table_name):
                columns.append({
                    "name": column["name"],
                    "type": str(column["type"]),
                    "nullable": column.get("nullable", True),
                    "default": column.get("default"),
                    "primary_key": column.get("primary_key", False),
                    "autoincrement": column.get("autoincrement", False)
                })
            
            # 获取主键
            primary_keys = inspector.get_pk_constraint(table_name).get("constrained_columns", [])
            
            # 获取外键
            foreign_keys = []
            for fk in inspector.get_foreign_keys(table_name):
                foreign_keys.append({
                    "name": fk["name"],
                    "constrained_columns": fk["constrained_columns"],
                    "referred_table": fk["referred_table"],
                    "referred_columns": fk["referred_columns"]
                })
            
            # 获取索引
            indexes = []
            for index in inspector.get_indexes(table_name):
                indexes.append({
                    "name": index["name"],
                    "columns": index["column_names"],
                    "unique": index.get("unique", False)
                })
            
            return {
                "table_name": table_name,
                "columns": columns,
                "primary_keys": primary_keys,
                "foreign_keys": foreign_keys,
                "indexes": indexes,
                "column_count": len(columns)
            }
            
        except Exception as e:
            logger.error(f"获取表结构失败: {str(e)}")
            raise DatabaseError(f"获取表结构失败: {str(e)}")
    
    @contextmanager
    def get_connection(self):
        """获取数据库连接上下文管理器"""
        if not self.is_connected or not self.engine:
            raise DatabaseError("数据库未连接")
        
        try:
            conn = self.engine.connect()
            yield conn
            conn.close()
        except Exception as e:
            logger.error(f"获取数据库连接失败: {str(e)}")
            raise DatabaseError(f"获取数据库连接失败: {str(e)}")
    
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()
            self.engine = None
            self.is_connected = False
            logger.info("数据库连接已关闭")
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "config": self.config,
            "connection_info": self.connection_info,
            "last_test_time": self.last_test_time.isoformat() if self.last_test_time else None,
            "is_connected": self.is_connected
        }


class DataSourceManager:
    """数据源管理器"""
    
    def __init__(self):
        self.connections: Dict[str, DataSourceConnection] = {}
    
    def add_connection(self, connection_id: str, config: Dict[str, Any]) -> DataSourceConnection:
        """添加数据源连接"""
        connection = DataSourceConnection(config)
        self.connections[connection_id] = connection
        return connection
    
    def get_connection(self, connection_id: str) -> Optional[DataSourceConnection]:
        """获取数据源连接"""
        return self.connections.get(connection_id)
    
    def remove_connection(self, connection_id: str):
        """移除数据源连接"""
        if connection_id in self.connections:
            self.connections[connection_id].close()
            del self.connections[connection_id]
    
    def test_all_connections(self) -> Dict[str, Dict[str, Any]]:
        """测试所有连接"""
        results = {}
        for connection_id, connection in self.connections.items():
            results[connection_id] = connection.test_connection()
        return results
    
    def get_connection_status(self) -> Dict[str, Dict[str, Any]]:
        """获取所有连接状态"""
        status = {}
        for connection_id, connection in self.connections.items():
            status[connection_id] = connection.to_dict()
        return status


# 全局数据源管理器
data_source_manager = DataSourceManager()


def create_data_source_connection(config: Dict[str, Any]) -> DataSourceConnection:
    """创建数据源连接的工厂函数"""
    return DataSourceConnection(config)


def get_supported_database_types() -> List[Dict[str, str]]:
    """获取支持的数据库类型列表"""
    return [
        {"code": DatabaseType.MYSQL, "name": "MySQL", "description": "MySQL数据库"},
        {"code": DatabaseType.POSTGRESQL, "name": "PostgreSQL", "description": "PostgreSQL数据库"},
        {"code": DatabaseType.SQLITE, "name": "SQLite", "description": "SQLite数据库"},
        {"code": DatabaseType.ORACLE, "name": "Oracle", "description": "Oracle数据库"},
        {"code": DatabaseType.SQLSERVER, "name": "SQL Server", "description": "Microsoft SQL Server"}
    ]